import pandas as pd
import paramiko
import time
import os
import re
from openpyxl import load_workbook

def read_server_info(excel_path):
    """
    从Excel文件中读取服务器连接信息
    :param excel_path: Excel文件路径
    :return: 包含服务器信息的DataFrame
    """
    try:
        df = pd.read_excel(excel_path)
        required_columns = ['IP地址', '端口', '账户', '密码', 'root账户', 'root密码']
        
        # 检查必要的列是否存在
        for col in required_columns:
            if col not in df.columns:
                print(f"错误: Excel文件中缺少必要的列 '{col}'")
                return None
        
        # 添加备注列（如果不存在）
        if '备注' not in df.columns:
            df['备注'] = ''
        
        return df
    except Exception as e:
        print(f"读取Excel文件时出错: {e}")
        return None

def connect_to_server(ip, port, username, password):
    """
    连接到服务器
    :param ip: 服务器IP地址
    :param port: 端口
    :param username: 用户名
    :param password: 密码
    :return: SSH客户端连接
    """
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(hostname=ip, port=port, username=username, password=password, timeout=10)
        return ssh
    except Exception as e:
        print(f"连接服务器 {ip}:{port} 失败: {e}")
        return None

def switch_to_root(ssh, root_username, root_password):
    """
    切换到root用户
    :param ssh: SSH连接
    :param root_username: root用户名
    :param root_password: root密码
    :return: 交互式shell通道
    """
    try:
        # 创建交互式shell
        channel = ssh.invoke_shell()
        time.sleep(1)
        
        # 执行su命令切换到root
        channel.send(f"su - {root_username}\n")
        time.sleep(1)
        
        # 发送root密码
        channel.send(f"{root_password}\n")
        time.sleep(2)
        
        # 检查是否成功切换到root
        channel.send("whoami\n")
        time.sleep(1)
        output = channel.recv(1024).decode('utf-8', errors='ignore')
        
        if root_username in output:
            print(f"成功切换到{root_username}用户")
            return channel
        else:
            print(f"切换到{root_username}用户失败")
            return None
    except Exception as e:
        print(f"切换到{root_username}用户时出错: {e}")
        return None

def get_java_processes(channel):
    """
    获取Java进程信息
    :param channel: 交互式shell通道
    :return: Java进程信息字符串
    """
    try:
        # 清空缓冲区
        if channel.recv_ready():
            channel.recv(1024)
        
        # 执行ps命令查找Java进程
        channel.send("ps -ef | grep java | grep -v grep\n")
        time.sleep(2)
        
        # 读取输出
        output = ""
        while channel.recv_ready():
            output += channel.recv(1024).decode('utf-8', errors='ignore')
        
        # 解析输出，提取进程号、包名和路径
        java_processes = []
        for line in output.splitlines():
            if "java" in line and not line.strip().startswith("grep"):
                # 尝试提取PID
                parts = line.split()
                if len(parts) > 1:
                    pid = parts[1]
                    
                    # 尝试提取jar包名或类名
                    jar_match = re.search(r'-jar\s+(\S+\.jar)', line)
                    class_match = re.search(r'(\w+\.)+\w+', line)
                    
                    if jar_match:
                        package = jar_match.group(1)
                    elif class_match:
                        package = class_match.group(0)
                    else:
                        package = "未知"
                    
                    # 尝试提取路径
                    path_match = re.search(r'(/\w+)+/\w+\.jar', line)
                    if path_match:
                        path = path_match.group(0)
                    else:
                        path = "未知"
                    
                    java_processes.append(f"PID: {pid}, 包名: {package}, 路径: {path}")
        
        return "\n".join(java_processes) if java_processes else "未找到Java进程"
    except Exception as e:
        print(f"获取Java进程信息时出错: {e}")
        return f"错误: {str(e)}"

def write_results_to_excel(excel_path, df):
    """
    将结果写入Excel文件
    :param excel_path: Excel文件路径
    :param df: 包含结果的DataFrame
    """
    try:
        # 方法1：使用新版pandas的ExcelWriter方式
        try:
            # 加载原始工作簿以保留格式
            book = load_workbook(excel_path)
            sheet_name = book.sheetnames[0]
            
            with pd.ExcelWriter(excel_path, engine='openpyxl', mode='a', if_sheet_exists='replace') as writer:
                # 删除现有的工作表
                if sheet_name in writer.sheets:
                    del writer.book[sheet_name]
                
                # 写入数据
                df.to_excel(writer, sheet_name=sheet_name, index=False)
            
            print(f"结果已成功写入到 {excel_path}")
            return
        except Exception as e1:
            print(f"使用方法1写入Excel文件时出错: {e1}")
        
        # 方法2：使用更简单的方式，但可能会丢失原始格式
        try:
            df.to_excel(excel_path, index=False)
            print(f"使用备用方法将结果写入到 {excel_path}")
            return
        except Exception as e2:
            print(f"使用方法2写入Excel文件时出错: {e2}")
        
        # 方法3：最后的尝试，使用openpyxl直接操作
        try:
            # 创建一个新的工作簿
            from openpyxl import Workbook
            wb = Workbook()
            ws = wb.active
            
            # 添加标题行
            for col_idx, column_name in enumerate(df.columns, 1):
                ws.cell(row=1, column=col_idx, value=column_name)
            
            # 添加数据行
            for row_idx, row in enumerate(df.values, 2):
                for col_idx, value in enumerate(row, 1):
                    ws.cell(row=row_idx, column=col_idx, value=value)
            
            # 保存工作簿
            wb.save(excel_path)
            print(f"使用openpyxl直接写入到 {excel_path}")
            return
        except Exception as e3:
            print(f"使用方法3写入Excel文件时出错: {e3}")
        
        raise Exception("所有写入方法都失败")
    except Exception as e:
        print(f"写入Excel文件时出错: {e}")
        print("无法保存结果到Excel文件。请检查文件是否被其他程序占用，或者尝试使用不同的文件名。")

def main():
    # 设置Excel文件路径
    excel_path = input("请输入Excel文件路径: ")
    
    if not os.path.exists(excel_path):
        print(f"错误: 文件 '{excel_path}' 不存在")
        return
    
    # 读取服务器信息
    server_df = read_server_info(excel_path)
    if server_df is None:
        return
    
    # 遍历每个服务器
    for index, row in server_df.iterrows():
        ip = row['IP地址']
        port = int(row['端口']) if not pd.isna(row['端口']) else 22
        username = row['账户']
        password = row['密码']
        root_username = row['root账户']
        root_password = row['root密码']
        
        print(f"\n处理服务器: {ip}:{port}")
        
        # 检查普通账户信息是否完整
        if pd.isna(username) or pd.isna(password):
            if pd.isna(root_username) or pd.isna(root_password):
                print(f"错误: 服务器 {ip} 的账户信息不完整")
                server_df.at[index, '备注'] = "账户信息不完整，无法连接"
                continue
            else:
                # 直接使用root账户
                username = root_username
                password = root_password
                print(f"使用root账户直接连接: {username}")
        
        # 连接服务器
        ssh = connect_to_server(ip, port, username, password)
        if ssh is None:
            server_df.at[index, '备注'] = "连接失败"
            continue
        
        # 如果使用的是普通账户且root账户信息完整，则切换到root
        if username != root_username and not pd.isna(root_username) and not pd.isna(root_password):
            channel = switch_to_root(ssh, root_username, root_password)
            if channel is None:
                ssh.close()
                server_df.at[index, '备注'] = "切换到root用户失败"
                continue
        else:
            # 直接使用当前连接创建通道
            channel = ssh.invoke_shell()
            time.sleep(1)
        
        # 获取Java进程信息
        java_info = get_java_processes(channel)
        
        # 更新DataFrame
        server_df.at[index, '备注'] = java_info
        
        # 关闭连接
        channel.close()
        ssh.close()
        print(f"服务器 {ip} 处理完成")
    
    # 将结果写入Excel
    # 创建一个新的Excel文件路径，避免覆盖原始文件
    result_excel_path = os.path.splitext(excel_path)[0] + "_results" + os.path.splitext(excel_path)[1]
    write_results_to_excel(result_excel_path, server_df)
    print(f"所有服务器处理完成，结果保存在: {result_excel_path}")

if __name__ == "__main__":
    main() 